# -*- coding:utf-8 -*-
import os, sys, re
import numpy as np
from metlib.kits import *
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch
from mpl_toolkits.basemap import Basemap
from matplotlib.cm import *
import matplotlib.patches as patches
from matplotlib.path import Path
from matplotlib.collections import PatchCollection
import matplotlib.collections
import mpl_toolkits.basemap.cm as bcm
from metlib.color import *
from metlib.plot.util import *
from pyproj import Proj
# from project.models import *
# from mapdata.models import *
from django.contrib.gis.geos import GEOSGeometry, Point, Polygon, LineString, MultiLineString, MultiPolygon
import gc
from metlib.color.cmap_utils import get_cmap as mm_get_cmap

# terrain2 = mm_get_cmap((terrain, (0.0, 0.75)))
# terrain2_r = mm_get_cmap((terrain_r, (0.25, 1.0)))
# cmap_d = {
#     'jet': jet,
#     'Spectral_r': Spectral_r,
#     'Spectral': Spectral,
#     'spectral_r': spectral,
#     'spectral': spectral_r,
#     'grey': Greys,
#     'grey_r': Greys_r,
#     'terrain': terrain,
#     'terrain_r': terrain_r,
#     'terrain2': terrain2,
#     'terrain2_r': terrain2_r,
#     # 功能性的:
#     'temperature': jet,
#     'solar': jet,
#     'rain': mm_get_cmap((spectral_r, (0.15, 1))),
#     'pressure': terrain_r,
#     'height': terrain
# }
myfont = matplotlib.font_manager.FontProperties(fname='%s/wqy-microhei.ttc' % DIRNAME(__file__))

def plot_lonlat_data(lons, lats, data,
                     fig_fname='fig.png', colorbar_fname=None,
                     target='map',
                     projection='merc',
                     pitem=None,
                     lon1=None, lon2=None, lat1=None, lat2=None,
                     vmin=None, vmax=None,
                     cmap='jet', level_n=64,
                     xpixels=None, pixel_per_grid=8, dpi=100,
                     min_xpixels=200, max_xpixels=5000,
                     title='', units='',
                     colorbar_title='',
                     ):
    """

    :param lons:
    :param lats:
    :param data:
    :param fig_fname:
    :param colorbar_fname:
    :param target: 'map', 'fig'
    :param projection:
    :param lon1:
    :param lon2:
    :param lat1:
    :param lat2:
    :param vmin:
    :param vmax:
    :param cmap:
    :param xpixels:
    :param pixel_per_grid:
    :param dpi:
    :return:
    """

    bbox = None
    boundary = None
    if pitem:
        try:
            pitemObj = WrasProjectItem.objects.get(pk=pitem)
            if pitemObj.item_type == 'Rect':
                bbox = pitemObj.geometry
                lon1, lat1, lon2, lat2 = bbox.extent
            elif pitemObj.item_type == 'MultiPolygon':
                lon1, lat1, lon2, lat2 = pitemObj.geometry.extent
                bbox = Polygon.from_bbox((lon1, lat1, lon2, lat2))
                boundary = pitemObj.geometry.coords
            else:
                pass  # 继续尝试用lon1, lat1 等...
        except WrasProjectItem.DoesNotExist:
            pass  # 继续尝试用lon1, lat1 等...

    # 清理lon1, lon2, lat1, lat2, lats, lons
    if lon1 is None:
        lon1 = np.min(lons)
    if lon2 is None:
        lon2 = np.max(lons)
    if lat1 is None:
        lat1 = np.min(lats)
    if lat2 is None:
        lat2 = np.max(lats)
    lon1 = float(lon1); lon2 = float(lon2); lat1 = float(lat1); lat2 = float(lat2)
    if lon1 > lon2:
        lon1, lon2 = lon2, lon1
    if lat1 > lat2:
        lat1, lat2 = lat2, lat1
    if len(np.shape(lats)) == 1 and len(np.shape(lons)) == 1:
        lons, lats = np.meshgrid(lons, lats)

    if bbox is None:
        bbox = Polygon.from_bbox((lon1, lat1, lon2, lat2))

    result = dict(lon1=lon1, lon2=lon2, lat1=lat1, lat2=lat2, pitem=pitem)

    # 清理vmin, vmax 等
    if vmin:
        vmin = float(vmin)
    else:
        vmin = np.nanmin(data)
    if vmax:
        vmax = float(vmax)
    else:
        vmax = np.nanmax(data)
    level_n = int(level_n)
    levels = np.linspace(vmin, vmax, level_n)
    if isinstance(cmap, (str, unicode)):
        cmap = cmap_d.get(cmap, jet)

    # resolution = None
    if target == 'map':
        resolution = None
    elif target == 'fig':
        lon_delta = lon2 - lon1
        if lon_delta >= 100:
            resolution = 'c'
        elif lon_delta >= 20:
            resolution = 'l'
        else:
            resolution = 'i'

    m = Basemap(projection=projection,
                llcrnrlat=lat1, urcrnrlat=lat2, llcrnrlon=lon1, urcrnrlon=lon2,
                resolution=resolution
                )
    x1, y1 = m(lon1, lat1)
    x2, y2 = m(lon2, lat2)
    if xpixels:
        xpixels = int(xpixels)
    else:
        xpixels = np.shape(data)[0] * pixel_per_grid

    xpixels = Limiter(min_xpixels, max_xpixels)(xpixels)

    if target == 'map':
        figscale = xpixels / (x2 - x1) / int(dpi)
    elif target == 'fig':
        figscale = 1000 / (x2 - x1) / int(dpi)

    x_size = (x2-x1) * figscale
    y_size = (y2-y1) * figscale
    if target == 'fig':
        x_size /= 0.8
        y_size /= 0.85

    fig = plt.figure(figsize=(x_size, y_size), facecolor='none', edgecolor='none', dpi=dpi)
    if target == 'fig':
        ax = fig.add_axes((0.07, 0.06, 0.8, 0.85))
        cax = fig.add_axes((0.91, 0.14, 0.02, 0.75))
        lax = fig.add_axes((0.88, 0.02, 0.11, 0.09), axisbg='none')
        lax.set_axis_off()
    else:
        ax = fig.add_axes((0, 0, 1, 1), axisbg='none')
        ax.set_axis_off()

    img = m.contourf(lons, lats, data, levels,
               ax=ax, latlon=True,
               extend='both',
               cmap=cmap)

    if boundary:
        paths = []
        for poly_arr in boundary:
            for subpoly_arr in poly_arr:
                lonlats_arr = np.array(subpoly_arr)
                x_arr, y_arr = m(lonlats_arr[:, 0], lonlats_arr[:, 1])
                xy = np.array([x_arr, y_arr]).transpose()
                edgecolor = 'none'
                if target == 'fig':
                    edgecolor = 'k'
                polygon = patches.Polygon(xy, closed=True, facecolor='none', edgecolor=edgecolor, linewidth=2.0)
                if target == 'fig':
                    ax.add_patch(polygon)
                paths.append(polygon.get_path())

        compound_path = Path.make_compound_path(*paths)

        for col in img.collections:
            col.set_clip_path(compound_path, ax.transData)

    transparent = True
    if target == 'fig':
        m.drawmeridians(easy_ticks(lon1, lon2), labels=[0, 0, 0, 1], ax=ax, color='0.7', linewidth=0.75)
        m.drawparallels(easy_ticks(lat1, lat2), labels=[1, 0, 0, 0], ax=ax, color='0.7', linewidth=0.75)
        m.drawcoastlines(ax=ax, linewidth=0.5, color='0.5')
        m.drawcountries(ax=ax, linewidth=0.5, color='0.5')
        m.readshapefile('%s/map/bou2_4p' % DIRNAME(__file__), 'bou2_4p', ax=ax, linewidth=0.5)

        # xian_mpolys = CommonPolygonMap.objects.filter(level='county', geometry__bboverlaps=bbox)
        # for xian_mpoly in xian_mpolys:
        #     for poly_arr in xian_mpoly.geometry.coords:
        #         for subpoly_arr in poly_arr:
        #             lonlats_arr = np.array(subpoly_arr)
        #             x_arr, y_arr = m(lonlats_arr[:, 0], lonlats_arr[:, 1])
        #             xy = np.array([x_arr, y_arr]).transpose()
        #             polygon = patches.Polygon(xy, closed=True, facecolor='none', edgecolor='r', linewidth=0.5)
        #             if target == 'fig':
        #                 ax.add_patch(polygon)

        # sheng_mpolys = CommonPolygonMap.objects.filter(level='province', geometry__bboverlaps=bbox)
        # for sheng_mpoly in sheng_mpolys:
        #     for poly_arr in sheng_mpoly.geometry.coords:
        #         for subpoly_arr in poly_arr:
        #             lonlats_arr = np.array(subpoly_arr)
        #             x_arr, y_arr = m(lonlats_arr[:, 0], lonlats_arr[:, 1])
        #             xy = np.array([x_arr, y_arr]).transpose()
        #             polygon = patches.Polygon(xy, closed=True, facecolor='none', edgecolor='k', linewidth=1.0)
        #             if target == 'fig':
        #                 ax.add_patch(polygon)

        ax.set_title(title, fontsize=16, fontproperties=myfont)
        cbar = plt.colorbar(img, cax=cax)
        cbar.set_ticks(easy_ticks(vmin, vmax))
        cbar_title = colorbar_title
        if units:
            cbar_title += u'  ( %s )' % units
        cbar.set_label(cbar_title, fontsize=12, fontproperties=myfont)
        # cbar.set_label(colorbar_title, fontsize=12)
        plt.setp(cax.get_yticklabels(), fontsize=11)
        transparent = False
        logo_fname = os.path.join(DIRNAME(__file__), 'logo.png')
        if (os.path.exists(logo_fname)):
            logo = plt.imread(logo_fname)
            lax.imshow(logo)

    fig.savefig(fig_fname, transparent=transparent, format='png')

    if target != 'fig' and colorbar_fname:
        cbarfig = plt.figure(figsize=(1, 4), facecolor='none', edgecolor='none', dpi=dpi)
        cax = cbarfig.add_axes((0.1, 0.05, 0.2, 0.9), zorder=20)
        chidax = cbarfig.add_axes((1.1, 1.1, 0.1, 0.1))
        cbackax = cbarfig.add_axes((0, 0, 1, 1), zorder=10)
        cbackax.set_axis_off()
        p_bbox = FancyBboxPatch((0, 0),10, 10, boxstyle="square,pad=0",ec='none',fc='w',alpha=0.4)
        cbackax.add_patch(p_bbox)

        cbar = plt.colorbar(img, cax=cax)
        cbar.set_ticks(easy_ticks(vmin, vmax))
        cbar_title = colorbar_title
        if units:
            cbar_title += u'  ( %s )' % units
        cbar.set_label(cbar_title, fontsize=12, fontproperties=myfont)
        # cbar.set_label(colorbar_title, fontsize=12)
        plt.setp(cax.get_yticklabels(), fontsize=10)

        cbarfig.savefig(colorbar_fname, transparent=transparent, format='png')

        cbarfig.clf()
        del cbarfig

    fig.clf()
    del fig
    del m
    plt.close()
    gc.collect()

    return result







